import sys
import time
import shutil
import subprocess
import hashlib
import ctypes
import os
import signal
import pefile
import psutil
from pathlib import Path
from typing import Optional


# === 日志初始化 ===
def init_logger():
    log_dir = Path(sys.executable).parent / "logs"
    log_dir.mkdir(parents=True, exist_ok=True)
    log_file = log_dir / "updater.log"

    import logging
    logger = logging.getLogger("Updater")
    logger.setLevel(logging.INFO)

    if logger.handlers:
        return logger

    formatter = logging.Formatter('%(asctime)s - %(levelname)s - [Updater] %(message)s')
    file_handler = logging.FileHandler(log_file, encoding='utf-8')
    file_handler.setFormatter(formatter)
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setFormatter(formatter)

    logger.addHandler(file_handler)
    logger.addHandler(console_handler)
    return logger


logger = init_logger()


# === 共享工具函数 ===
def verify_sha256(file_path: Path, expected_hash: str) -> bool:
    if not file_path.exists():
        logger.error(f"[校验] 文件不存在：{file_path}")
        return False
    sha256 = hashlib.sha256()
    try:
        with open(file_path, 'rb') as f:
            for chunk in iter(lambda: f.read(8192), b""):
                sha256.update(chunk)
        actual = sha256.hexdigest().lower()
        expected = expected_hash.lower()
        logger.info(f"[校验] 实际 SHA256: {actual}")
        logger.info(f"[校验] 期望 SHA256: {expected}")
        return actual == expected
    except Exception as e:
        logger.error(f"[校验] 校验失败: {e}", exc_info=True)
        return False


def is_64bit_file(file_path: Path) -> Optional[bool]:
    logger.info(f"[架构识别] 检测文件架构: {file_path}")
    if not file_path.exists() or not file_path.is_file():
        logger.error(f"[架构识别] 文件不存在或不是文件: {file_path}")
        return None
    try:
        pe = pefile.PE(str(file_path))
        if pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_AMD64"]:
            logger.info(f"[架构识别] {file_path} 是64位文件")
            return True
        elif pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_I386"]:
            logger.info(f"[架构识别] {file_path} 是32位文件")
            return False
        else:
            logger.warning(f"[架构识别] {file_path} 是未知架构")
            return None
    except Exception as e:
        logger.error(f"[架构识别] 解析 {file_path} 失败: {e}", exc_info=True)
        return None


def is_admin() -> bool:
    try:
        result = ctypes.windll.shell32.IsUserAnAdmin()
        logger.info(f"[权限检查] 管理员权限: {'是' if result else '否'}")
        return result
    except Exception as e:
        logger.error(f"[权限检查] 检查失败: {e}", exc_info=True)
        return False


# === 进程清理工具函数 ===
def kill_process_by_name(process_name: str) -> bool:
    """强制关闭指定名称的进程（解决文件占用问题）"""
    try:
        killed = False
        for proc in psutil.process_iter(['name', 'pid']):
            if proc.info['name'] == process_name:
                logger.warning(f"[进程清理] 发现残留进程 {process_name} (PID: {proc.info['pid']})，强制关闭")
                proc.kill()
                killed = True
                time.sleep(0.5)  # 等待资源释放
        return killed
    except Exception as e:
        logger.error(f"[进程清理] 关闭进程失败: {e}", exc_info=True)
        return False


# === 文件操作工具函数 ===
def force_delete(path: Path) -> bool:
    """强制删除文件/目录（处理占用情况）"""
    logger.info(f"[删除] 尝试删除: {path}")
    if not path.exists():
        logger.info(f"[删除] 已不存在，无需处理")
        return True
    try:
        if path.is_file():
            path.unlink(missing_ok=True)
            if path.exists():
                # 标记为重启后删除
                delete_path = str(path) + ".delete_on_reboot"
                os.rename(path, delete_path)
                logger.warning(f"[删除] 文件被占用，标记为重启后删除: {delete_path}")
                return False
            logger.info(f"[删除] 文件删除成功: {path}")
            return True
        else:
            shutil.rmtree(path, ignore_errors=True)
            if path.exists():
                logger.warning(f"[删除] 目录仍存在（可能被占用）: {path}")
                return False
            logger.info(f"[删除] 目录删除成功: {path}")
            return True
    except Exception as e:
        logger.error(f"[删除] {path} 删除失败: {e}", exc_info=True)
        return False


def copy_with_retry(src: Path, dst: Path, max_retries: int = 3) -> bool:
    """带重试机制的文件复制（处理临时占用）"""
    logger.info(f"[复制] 开始复制: {src} -> {dst}（最多重试{max_retries}次）")
    for i in range(max_retries):
        try:
            if src.is_dir():
                shutil.copytree(src, dst, dirs_exist_ok=True)
                logger.info(f"[复制] 目录复制完成: {src} -> {dst}")
            else:
                dst.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(src, dst)
                if not dst.exists() or dst.stat().st_size != src.stat().st_size:
                    raise Exception(
                        f"文件大小不匹配（源: {src.stat().st_size}，目标: {dst.stat().st_size if dst.exists() else 0}）")
                logger.info(f"[复制] 文件复制完成: {src} -> {dst}")
            return True
        except Exception as e:
            logger.warning(f"[复制] 第 {i + 1} 次复制失败: {e}")
            if i < max_retries - 1:
                time.sleep(1)
    logger.error(f"[复制] 超过最大重试次数，{src} -> {dst} 复制失败")
    return False


# === 版本文件更新 ===
def update_version_file(extracted_ver_file: Path, target_dir: Path) -> bool:
    """更新版本文件（确保在所有文件替换完成后执行）"""
    logger.info(f"[版本更新核心] 开始更新版本文件: {extracted_ver_file} -> {target_dir}/version.txt")
    try:
        if not extracted_ver_file.exists():
            raise Exception(f"新版本文件不存在: {extracted_ver_file}")

        # 读取新版本号
        with open(extracted_ver_file, 'r', encoding='utf-8') as f:
            new_version = f.read().strip()
        logger.info(f"[版本更新核心] 新版本号: {new_version}")

        target_ver_file = target_dir / "version.txt"
        # 先删除旧版本文件
        if target_ver_file.exists() and not force_delete(target_ver_file):
            raise Exception(f"无法删除旧版本文件: {target_ver_file}")

        # 复制新版本文件
        if not copy_with_retry(extracted_ver_file, target_ver_file):
            raise Exception(f"无法复制版本文件到目标位置: {target_ver_file}")

        # 校验版本内容
        with open(target_ver_file, 'r', encoding='utf-8') as f:
            updated_ver = f.read().strip()
        if updated_ver != new_version:
            raise Exception(f"版本文件内容不匹配（预期: {new_version}，实际: {updated_ver}）")

        logger.info(f"[版本更新核心] 版本文件更新成功: {target_ver_file}")
        return True
    except Exception as e:
        logger.error(f"[版本更新核心] 版本文件更新失败: {e}", exc_info=True)
        return False


# === 更新核心逻辑 ===
def replace_main_executable(extracted_exe: Path, target_exe: Path, expected_sha256: str) -> bool:
    """替换主程序可执行文件"""
    logger.info(f"[主程序替换核心] 开始替换主程序: {extracted_exe} -> {target_exe}")

    # 1. 校验更新文件完整性
    if not verify_sha256(extracted_exe, expected_sha256):
        logger.error("[主程序替换核心] 更新文件SHA256校验失败，可能被篡改")
        return False

    # 2. 校验架构兼容性
    target_is_64 = is_64bit_file(target_exe) if target_exe.exists() else None
    extracted_is_64 = is_64bit_file(extracted_exe)
    if target_is_64 is not None and extracted_is_64 is not None and target_is_64 != extracted_is_64:
        logger.error(
            f"[主程序替换核心] 架构不兼容：目标文件是{'64' if target_is_64 else '32'}位，更新文件是{'64' if extracted_is_64 else '32'}位")
        return False

    # 3. 处理目标文件占用
    if target_exe.exists() and not force_delete(target_exe):
        logger.error(f"[主程序替换核心] 目标文件 {target_exe} 被占用，无法删除")
        return False

    # 4. 复制新文件到目标位置
    if not copy_with_retry(extracted_exe, target_exe):
        logger.error(f"[主程序替换核心] 无法将更新文件复制到目标位置 {target_exe}")
        return False

    # 5. 最终校验新文件
    if not verify_sha256(target_exe, expected_sha256):
        logger.error("[主程序替换核心] 替换后文件校验失败，更新损坏")
        if target_exe.exists():
            force_delete(target_exe)
        return False

    logger.info(f"[主程序替换核心] 主程序替换成功: {target_exe}")
    return True


def copy_all_files(extracted_dir: Path, target_dir: Path) -> bool:
    """复制所有更新文件（除主程序外）"""
    logger.info(f"[文件复制核心] 开始复制所有文件: {extracted_dir} -> {target_dir}")

    main_exe_name = Path(sys.argv[2]).name  # 主程序文件名
    updater_exe_name = Path(sys.executable).name  # 更新器自身文件名
    pending_updater = None
    pending_files = []  # 延迟替换文件列表

    for item in extracted_dir.rglob("*"):
        relative_path = item.relative_to(extracted_dir)
        target_path = target_dir / relative_path

        # 跳过主程序（单独处理）
        if item.is_file() and item.name == main_exe_name:
            logger.info(f"[文件复制] 跳过主程序文件: {item}")
            continue

        # 暂存更新器自身（延迟处理）
        if item.is_file() and item.name == updater_exe_name:
            pending_updater = (item, target_path)
            continue

        # 确保目标目录存在
        if item.is_dir():
            target_path.mkdir(parents=True, exist_ok=True)
            continue

        # 复制文件，失败则加入延迟列表
        if not copy_with_retry(item, target_path, max_retries=3):
            logger.warning(f"[文件复制] {item} 暂时无法复制，加入延迟替换列表")
            pending_files.append((item, target_path))

    # 处理延迟替换文件（通过批处理）
    if pending_files and not handle_pending_files(pending_files):
        return False

    # 处理更新器自身替换
    if pending_updater and not handle_updater_update(pending_updater[0], pending_updater[1]):
        return False

    logger.info("[文件复制核心] 所有文件复制完成（包括延迟替换）")
    return True


def handle_pending_files(pending_files: list) -> bool:
    """通过批处理延迟替换被占用的文件"""
    try:
        temp_dir = Path(sys.executable).parent / "update_pending"
        temp_dir.mkdir(exist_ok=True)

        batch_commands = []
        for src, dst in pending_files:
            temp_file = temp_dir / src.name
            if not copy_with_retry(src, temp_file, max_retries=1):
                logger.error(f"[延迟替换] 无法复制 {src} 到临时目录，放弃该文件")
                continue
            batch_commands.append(f'echo 替换 {dst.name}...')
            batch_commands.append(f'del /f /q "{dst}" 2>nul')
            batch_commands.append(f'copy /y "{temp_file}" "{dst}" 2>nul')

        # 生成批处理脚本
        batch_script = temp_dir / "replace_pending.bat"
        with open(batch_script, "w", encoding="utf-8") as f:
            f.write("@echo off\n")
            f.write("echo 正在处理延迟替换的文件...\n")
            f.write("timeout /t 3 /nobreak >nul\n")  # 等待主程序退出
            f.write("\n".join(batch_commands))
            f.write(f'\ndel /f /q "{batch_script}"\n')  # 清理自身
            f.write("exit\n")

        # 启动批处理
        subprocess.Popen(
            str(batch_script),
            shell=True,
            close_fds=True,
            creationflags=subprocess.CREATE_NEW_PROCESS_GROUP
        )
        logger.info(f"[延迟替换核心] 启动批处理处理 {len(pending_files)} 个文件")
        return True
    except Exception as e:
        logger.error(f"[延迟替换核心] 处理失败: {e}", exc_info=True)
        return False


def handle_updater_update(src_updater: Path, dst_updater: Path) -> bool:
    """通过批处理延迟替换更新器自身"""
    try:
        # 复制为临时文件
        temp_updater = dst_updater.with_name(f"{dst_updater.stem}_temp{dst_updater.suffix}")
        if not copy_with_retry(src_updater, temp_updater):
            raise Exception("无法复制更新器到临时文件")

        # 创建批处理脚本
        batch_script = dst_updater.parent / "update_updater.bat"
        with open(batch_script, "w", encoding="utf-8") as f:
            f.write(f"""
@echo off
echo 正在更新更新器...
timeout /t 2 /nobreak >nul  # 等待原更新器退出
if exist "{dst_updater}" (
    del /f /q "{dst_updater}"
    if errorlevel 1 (
        echo 删除旧更新器失败，将在重启后替换
        ren "{temp_updater}" "{dst_updater.name}"
        goto end
    )
)
ren "{temp_updater}" "{dst_updater.name}"
echo 更新器更新完成
:end
del /f /q "{batch_script}"
exit
            """)

        # 启动批处理
        subprocess.Popen(
            str(batch_script),
            shell=True,
            close_fds=True,
            creationflags=subprocess.CREATE_NEW_PROCESS_GROUP
        )
        logger.info(f"[更新器替换核心] 启动批处理替换更新器: {batch_script}")
        return True
    except Exception as e:
        logger.error(f"[更新器替换核心] 处理失败: {e}", exc_info=True)
        return False


def cleanup_temp_files(temp_dir: Path) -> None:
    """清理临时更新文件"""
    logger.info(f"[清理核心] 开始清理临时目录: {temp_dir}")
    try:
        if temp_dir.exists():
            shutil.rmtree(temp_dir, ignore_errors=True)
            if not temp_dir.exists():
                logger.info(f"[清理核心] 临时目录清理成功: {temp_dir}")
            else:
                logger.warning(f"[清理核心] 临时目录仍存在（可能被占用）: {temp_dir}")
        else:
            logger.info(f"[清理核心] 临时目录不存在，无需清理")
    except Exception as e:
        logger.warning(f"[清理核心] 清理失败: {e}", exc_info=True)


def restart_application(target_exe: Path) -> None:
    """重启应用程序"""
    logger.info(f"[重启核心] 尝试重启应用程序: {target_exe}")
    try:
        # 使用cmd启动避免更新器窗口阻塞
        subprocess.Popen(
            f'cmd /c "start "" "{target_exe}" "',
            shell=True,
            close_fds=True,
            creationflags=subprocess.CREATE_NEW_PROCESS_GROUP
        )
        logger.info("[重启核心] 应用程序启动命令已发送")
    except Exception as e:
        logger.error(f"[重启核心] 重启失败: {e}", exc_info=True)
        ctypes.windll.user32.MessageBoxW(
            0,
            f"更新完成，但自动重启失败。\n请手动启动程序：\n{target_exe}",
            "更新完成",
            0x40 | 0x1  # 信息图标 + 确定按钮
        )


# === 主函数 ===
def main():
    logger.info("===== 更新器启动 =====")
    logger.info(f"[启动核心] 命令行参数: {sys.argv}")

    # 校验参数（共5个参数：updater.exe <解压后主程序> <目标主程序> <SHA256> <版本文件>）
    if len(sys.argv) != 5:
        error_msg = "参数错误！正确格式: updater.exe <解压后的主程序路径> <目标主程序路径> <预期SHA256> <新版本文件路径>"
        logger.error(f"[启动核心] {error_msg}")
        ctypes.windll.user32.MessageBoxW(0, error_msg, "更新失败", 0x10)  # 错误图标
        return

    # 解析参数
    try:
        extracted_exe_path = Path(sys.argv[1]).resolve()
        target_exe_path = Path(sys.argv[2]).resolve()
        expected_sha256 = sys.argv[3]
        extracted_ver_file = Path(sys.argv[4]).resolve()  # 新版本文件路径
    except Exception as e:
        error_msg = f"解析参数失败: {str(e)}"
        logger.error(f"[启动核心] {error_msg}", exc_info=True)
        ctypes.windll.user32.MessageBoxW(0, error_msg, "更新失败", 0x10)
        return

    # 校验路径有效性
    if not extracted_exe_path.exists():
        error_msg = f"解压后的主程序不存在: {extracted_exe_path}"
        logger.error(f"[启动核心] {error_msg}")
        ctypes.windll.user32.MessageBoxW(0, error_msg, "更新失败", 0x10)
        return
    if not extracted_ver_file.exists():
        error_msg = f"新版本文件不存在: {extracted_ver_file}"
        logger.error(f"[启动核心] {error_msg}")
        ctypes.windll.user32.MessageBoxW(0, error_msg, "更新失败", 0x10)
        return

    # 校验管理员权限
    if not is_admin():
        error_msg = "更新需要管理员权限。\n请右键程序选择“以管理员身份运行”"
        logger.error(f"[启动核心] {error_msg}")
        ctypes.windll.user32.MessageBoxW(0, error_msg, "权限不足", 0x10)
        return

    try:
        # 强制关闭主程序进程（解决文件占用）
        main_process_name = target_exe_path.name
        logger.info(f"[启动核心] 尝试关闭主程序进程: {main_process_name}")
        kill_process_by_name(main_process_name)
        time.sleep(1)  # 等待资源释放

        # 获取目录路径
        extracted_dir = extracted_exe_path.parent
        target_dir = target_exe_path.parent

        # 1. 复制所有文件（除主程序外）
        logger.info("[更新流程核心] 开始执行文件复制阶段")
        if not copy_all_files(extracted_dir, target_dir):
            raise Exception("文件复制失败，更新终止")

        # 2. 替换主程序
        logger.info("[更新流程核心] 开始执行主程序替换阶段")
        if not replace_main_executable(extracted_exe_path, target_exe_path, expected_sha256):
            raise Exception("主程序替换失败，更新终止")

        # 3. 最后更新版本文件（确保文件替换完成后再更新版本号）
        logger.info("[更新流程核心] 开始执行版本文件更新阶段")
        if not update_version_file(extracted_ver_file, target_dir):
            raise Exception("版本文件更新失败，更新终止")

        # 4. 清理临时文件
        logger.info("[更新流程核心] 开始执行清理阶段")
        cleanup_temp_files(extracted_dir.parent)

        # 5. 自动重启应用
        logger.info("[更新流程核心] 开始执行重启阶段")
        restart_application(target_exe_path)

        logger.info("===== 更新成功 =====")

    except Exception as e:
        error_msg = f"更新失败: {str(e)}\n请查看 logs/updater.log 获取详情"
        logger.error(f"[更新流程核心] {error_msg}", exc_info=True)
        ctypes.windll.user32.MessageBoxW(
            0,
            error_msg,
            "更新失败",
            0x10  # 错误图标
        )


if __name__ == "__main__":
    main()